import re
import os
import logging
from lxml import etree
from neuro.base import Object
from neuro.exceptions import FileNotFoundException
from neuro.exceptions import IllegalArgumentException
from neuro.exceptions import BaseException
import neuro.strings as strings
import neuro.filesystem as filesystem
from neuro.common import isInteger, isFloat
import neuro.config as config

logger = logging.getLogger('neuro.apps.freesurfer')

class FreeSurfer(Object):
    '''
    FreeSurfer factory
    '''
    @staticmethod
    def create(output_dir, version):
        '''
        Factory method for creating FreeSurfer instances given a version ::
            
            >>> FreeSurfer.create("4.5.0")
            <neuro.apps.xnat.FreeSurfer_450 object at 0x1621ad0>
            
        :param version: FreeSurfer version 
        :type version: str
        :returns: FreeSurfer instance
        :rtype: :class:`FreeSurfer`
        '''
        if(version == '4.5.0'):
            return FreeSurfer_450(output_dir)
        else:
            raise UnknownFreeSurferVersion(version)

class FreeSurfer_450(Object):
    '''
    FreeSurfer 4.5.0
    '''
    def __init__(self, output_dir):
        '''
        Constructor
        '''
        self.output_dir = output_dir
    
    def genreport(self):
        '''
        Generate FreeSurfer report file
        
        :returns: XML
        :rtype: lxml.etree.Element 
        '''
        aseg = os.path.join(self.output_dir, 'stats', 'aseg.stats')
        lh_aparc = os.path.join(self.output_dir, 'stats', 'lh.aparc.stats')
        rh_aparc = os.path.join(self.output_dir, 'stats', 'rh.aparc.stats')
        
        aseg_xml = self._stats_to_xml(aseg, 'aseg_stats')
        lh_aparc_xml = self._stats_to_xml(lh_aparc, 'lh_aparc_stats')
        rh_aparc_xml = self._stats_to_xml(rh_aparc, 'rh_aparc_stats')
        
        derivs = etree.Element('derivs')
        derivs.append(aseg_xml.xpath('/derivs/*')[0])
        derivs.append(lh_aparc_xml.xpath('/derivs/*')[0])
        derivs.append(rh_aparc_xml.xpath('/derivs/*')[0])
        
        return derivs
        
    def _stats_to_xml(self, resource, group):
        '''
        Serialize stats file to XML
        
        :param resource: Stats filename or file-like object
        :type resource: str|file-like object
        :param group: Group label
        :type group: str
        :returns: lxml Element
        :rtype: lxml.etree.Element
        '''
        if isinstance(resource, basestring):
            ## --- change to 'rU' if newlines become an issue across platforms
            resource = open(resource, 'rb') 
        
        contents = resource.read().split('\n')
        if len(contents) == 0:
            raise StatsToXMLException(StatsToXMLException.RESOURCE_EMPTY, resource)
        
        ## --- initialize
        headers = []
        payload = []
    
        ## --- store headers and payload into separate arrays
        for row in contents:
            if(row != ""):
                if(row[0] == "#"):
                    headers.append(row)
                else:
                    payload.append(row)
                    
        ## --- create the beginning of the xml tree
        root = etree.Element("derivs")
        group_node = etree.Element("group", label=group)
        root.append(group_node)
        header = etree.Element("group", label="headers")
        group_node.append(header)
    
        ## --- initialize
        regs = []
        columns = []
    
        ## --- xml-ify all header fields of interest (hold your breath)
        for row in headers:
            deriv = None
            _text = None
            _label = None
            _units = None
    
            ## --- special regular expressions, add here if header line is oddly formatted.
            ## --- make sure to add special handler below
            version = strings.regex("# cmdline.*freesurfer-(.*)/", row)
            maskVox = strings.regex("# Measure BrainMask, (BrainMaskNVox),.*,(.*),(.*)", row)
            segVoxN = strings.regex("# Measure BrainSeg, (BrainSegNVox),.*,(.*),(.*)", row)
            lhwm = strings.regex("# surface-based-volume .* lh-cerebral-white-matter (.*)", row)
            rhwm = strings.regex("# surface-based-volume .* rh-cerebral-white-matter (.*)", row)
            numVert = strings.regex("# Measure Cortex, (NumVert),.*,(.*),(.*)", row)
            columns = strings.regex("(# ColHeaders.*)", row)
    
            ## --- standard regular expressions, usually you will add additional fields here
            if(len(regs) == 0):
                regs.append("# Measure Cortex, (SurfArea),.*,(.*),(.*)")
                regs.append("# Measure BrainMask, (BrainMaskVol),.*,(.*),(.*)")
                regs.append("# Measure BrainSegNotVent, (BrainSegVolNotVent),.*,(.*),(.*)")
                regs.append("# Measure BrainSeg, (BrainSegVol),.*,(.*),(.*)")
                regs.append("# Measure IntraCranialVol, (ICV),.*,(.*),(.*)")
    
            for reg in regs:
                deriv = strings.regex(reg, row)
                if(deriv):
                    break
                    
            ## --- handlers for "deriv" and other special header fields
            if(deriv):
                _label = deriv[0].strip()
                _text = deriv[1].strip()
                _units = deriv[2].lower().strip()
            if(version):
                _label = "version"
                _text = version[0].strip()
                _units = ""
            if(maskVox):
                _label = maskVox[0].strip()
                _text = maskVox[1].strip()
                _units = "N"
            if(segVoxN):
                _label = segVoxN[0].strip()
                _text = segVoxN[1].strip()
                _units = "N"
            if(lhwm):
                _label = "L_Surf_Cbral_wm"
                _text = lhwm[0].strip()
                _units = "mm^3"
            if(rhwm):
                _label = "R_Surf_Cbral_wm"
                _text = rhwm[0].strip()
                _units = "mm^3"
            if(numVert):
                _label = numVert[0].strip()
                _text = numVert[1].strip()
                _units = "N"
    
            ## --- add new derivative to header group_node
            if(_text != None):
                element = etree.Element("deriv", units=_units.lower(), label=_label.lower())
                element.text = _text
                header.append(element)
    
        ## --- quick little sanity check, we need header columns for zip
        if(not columns):
            raise StatsToXMLException(StatsToXMLException.MISSING_COLUMN_HEADERS, resource)
    
        ## --- do some explosion
        columns = columns[0]
        columns = strings.explode(columns.strip(), " ", strip=True, collapse=True)
        columns = columns[2::]
    
        ## --- more sanity checking, are we where we should be?
        if(columns[0] != "Index" and columns[0] != "StructName"):
            raise StatsToXMLException(StatsToXMLException.FORMAT_ERROR, resource)
                    
        ## --- process the payload
        for row in payload:
            ## --- explode the payload row into an array by spaces
            row = strings.explode(row.strip(), " ", strip=True, collapse=True)
            row = dict(zip(columns, row)) ## --- blissful memory hungry op
    
            ## --- create a new group_node labeled by the structure name
            struct_label = row["StructName"].replace("-", "_")
            structElement = etree.Element("group", label=struct_label.lower())
            group_node.append(structElement)
    
            ## --- we're going to iterate, so let's remove this rather than make a condition later
            del(row["StructName"])
    
            ## --- iterate
            for key,value in row.items():
                ## --- create new derivative and add it to the structure group and add to parent group
                units = ""
                if isInteger(value):
                    units = "int"
                elif isFloat(value):
                    units = "float"
                deriv_label = key
                derivElement = etree.Element("deriv", label=deriv_label.lower(), units=units)
                derivElement.text = value
                structElement.append(derivElement)
        
        return root

class UnknownFreeSurferVersion(BaseException):
    '''
    Unknown FreeSurfer version
    '''
    def __init__(self, version):
        '''
        Constructor
        
        :param version: FreeSurfer version
        :type version: str
        '''
        BaseException.__init__(self)
        self._version = version
        
    def getMessage(self):
        '''
        Get custom message

        :rtype: str
        '''
        return "Unknown FreeSurfer version: " + self._version
    
class StatsToXMLException(BaseException):
    '''
    FreeSurfer stats to XML exception
    '''
    RESOURCE_EMPTY=0
    MISSING_COLUMN_HEADERS=1
    FORMAT_ERROR=2
    
    def __init__(self, type, resource):
        '''
        Constructor

        :param type: Type of exception
        :type type: int
        :param resource: Resource
        :type filename: str|file-like-object
        '''
        BaseException.__init__(self)
        
        self._type = type
        self._resource = resource

    def getMessage(self):
        '''
        Get custom message

        :rtype: str
        '''
        if self._type == StatsToXMLException.RESOURCE_EMPTY:
            return "Resource is empty: " + str(self._resource)
        elif self._type == StatsToXMLException.MISSING_COLUMN_HEADERS:
            return "Unable to find column headers (# ColHeaders ...): " + str(self._resource)
        elif self._type == StatsToXMLException.FORMAT_ERROR:
            return "Format error. Expecting first column header to be \"Index\" or \"StructName\": " + str(self._resource)
